import numpy as np
import gradio as gr
from PIL import Image
from pinecone import Pinecone
from collections import Counter

# Pinecone 配置
PINECONE_API_KEY = "f83eac6c-0929-44a6-bd90-6f48dba6ea56"
PINECONE_INDEX = "mnist-index"

# 初始化 Pinecone 客户端
pinecone_client = Pinecone(api_key=PINECONE_API_KEY)
vector_db = pinecone_client.Index(PINECONE_INDEX)

def normalize_and_flatten(img_array):
    """归一化图像数组并展平"""
    return (np.array(img_array) / 255.0 * 16).flatten()

def prepare_image(raw_input):
    """预处理输入图像"""
    if raw_input is None or not isinstance(raw_input, np.ndarray):
        return None

    # 转换为灰度图像
    gray_image = np.mean(raw_input, axis=2).astype(np.uint8) if raw_input.ndim == 3 else raw_input
    
    # 调整大小并转换为 PIL Image
    resized_image = Image.fromarray(gray_image, "L").resize((8, 8))
    
    return normalize_and_flatten(resized_image)

def get_nearest_neighbors(vector, k=11):
    """从 Pinecone 获取最近邻"""
    try:
        response = vector_db.query(vector=vector.tolist(), top_k=k, include_metadata=True)
        return response['matches']
    except Exception as e:
        print(f"查询 Pinecone 时出错: {e}")
        return []

def majority_vote(labels):
    """执行多数投票"""
    return Counter(labels).most_common(1)[0][0] if labels else None

def digit_recognition(sketch_input):
    """主要的数字识别函数"""
    processed_vector = prepare_image(sketch_input)
    if processed_vector is None:
        return "无效输入"

    neighbors = get_nearest_neighbors(processed_vector)
    if not neighbors:
        return "无法找到相似的数字"

    neighbor_labels = [int(match['metadata']['label']) for match in neighbors]
    predicted_digit = majority_vote(neighbor_labels)

    return int(predicted_digit) if predicted_digit is not None else "无法确定数字"

# Gradio 界面配置
interface = gr.Interface(
    fn=digit_recognition,
    inputs=gr.Sketchpad(label="在此绘制数字"),
    outputs=gr.Textbox(label="识别结果"),
    title="手写数字识别系统",
    description="请在画板上绘制一个0到9之间的数字，系统将尝试识别它。"
)

# 启动 Gradio 应用
if __name__ == "__main__":
    interface.launch(share=True)